package handler

import (
	"fmt"
	"github.com/fatih/color"
)

func consumer(ips []string) {
	for _, ip := range ips {
		ipChan <- ip
	}
	close(ipChan)
}

func producer(fn func(ip string) (stdout, stderr []byte, err error), printStatuses []string, printFields []string) {
	for ip := range ipChan {
		res := newResult(ip, printStatuses, printFields)
		stdout, stderr, err := fn(ip)
		if stdout != nil {
			res.stdout = stdout
		}
		if stderr != nil {
			res.stderr = stderr
		}
		if err != nil {
			res.stderr = append(res.stderr, []byte(err.Error())...)
		}
		if err == nil {
			res.status = true
		}
		resChan <- res
		wg.Done()
	}
}

func readFromRes() {
	for {
		if res, ok := <-resChan; ok {

			if res.printFailed && !res.status {
				if res.printIP {
					fmt.Println(res.ip)
				}
				if res.printStdout && len(res.stdout) > 0 {
					color.Green(string(res.stdout))
				}
				if res.printStderr && len(res.stderr) > 0 {
					color.Red(string(res.stderr))
				}
			}
			if res.printSucceed && res.status {
				if res.printIP {
					fmt.Println(res.ip)
				}
				if res.printStdout && len(res.stdout) > 0 {
					color.Green(string(res.stdout))
				}
				if res.printStderr && len(res.stderr) > 0 {
					color.Yellow(string(res.stderr))
				}
			}
			wg.Done()
			continue
		}
		break
	}
}

func Handler(fork int, ips []string, printStatuses []string, printFields []string, fn func(ip string) (stdout, stderr []byte, err error)) {
	initChan(fork)

	go consumer(ips)
	wg.Add(len(ips) * 2)
	for i := 0; i < fork; i++ {
		go producer(fn, printStatuses, printFields)
	}
	go readFromRes()
	wg.Wait()
	close(resChan)
}
